#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import numpy as np
import pandas as pd
from joblib import wrap_non_picklable_objects

from gplearn_evolve.base.utils import d_type_exam
# from gplearn_evolve.base.utils import d_map
from gplearn_evolve.base.utils import d_map_limit

from gplearn_evolve.base.utils import keys_core_pd
from gplearn_evolve.base.utils import keys_core_pd_v2

from gplearn_evolve.base.config import cw_low_limit
from gplearn_evolve.base.config import cw_up_limit

from statsmodels.regression.rolling import RollingOLS
import statsmodels.api as sm

# 1
@wrap_non_picklable_objects
def def_keyfunc_cw_delay_qoq(keys, key_by_path=True):
    def gp_cw_delay_qoq(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = 1

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).shift(params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_cw_delay_qoq

# 2
def def_keyfunc_ts_cw_delay(keys, key_by_path=True):
    def gp_ts_cw_delay(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map_limit(d_type_exam(d), base=2, lowlimit=1, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).shift(params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_delay

# 3
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_corr(keys, key_by_path=True):
    def gp_ts_cw_corr(X, Y, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X, Y)
        else:
            values = keys_core_pd_v2(load_keys, X, Y)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.iloc[:,0].rolling(window=params,
                                                                                 min_periods=params).corr(x.iloc[:,1]))

        # TODO - 将inf设置为nan
        # values[np.isinf(values['f1'])] = np.nan
        values[np.isinf(values)] = np.nan

        del load_keys

        return values.values

    return gp_ts_cw_corr

# 4
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_min(keys, key_by_path=True):
    def gp_ts_cw_min(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                       min_periods=params).min())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_min

# 5
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_max(keys, key_by_path=True):
    def gp_ts_cw_max(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params).max())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_max

# 6
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_std(keys, key_by_path=True):
    def gp_ts_cw_std(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params).std())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_std

# 7
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_sum(keys, key_by_path=True):
    def gp_ts_cw_sum(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                       min_periods=params).sum())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_sum

# 8
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_norm(keys, key_by_path=True):
    def gp_ts_cw_norm(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values_mean = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params).mean())

        values_std = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params).std())

        values = (values - values_mean) / values_std

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_norm


'''rolling ols'''
# 9
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ols_residual(keys, key_by_path=True):
    def gp_ts_cw_ols_residual(X, Y, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X, Y)
        else:
            values = keys_core_pd_v2(load_keys, X, Y)

        def rolling_ols(df, params):
            try:
                endog = df.iloc[:, 0]
                exog = sm.add_constant(df.iloc[:, 1])
                mod = RollingOLS(endog=endog,
                                 exog=exog,
                                 window=params,
                                 min_nobs=params,
                                 missing='skip')

                model_res = mod.fit(params_only=True)

                df['f1'] = endog - (model_res.params * exog).sum(axis=1)

                return df['f1']
            except:
                ret_df = df.copy()
                ret_df.iloc[:, :] = np.nan

                return ret_df['f1']

        res = values.groupby(as_index=False,
                             level='stock_code',
                             sort=False).apply(rolling_ols, params)

        # TODO - 将inf设置为nan
        res[np.isinf(res)] = np.nan

        del load_keys

        return res.values

    return gp_ts_cw_ols_residual

# 10
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ols_rsq(keys, key_by_path=True):
    def gp_ts_cw_ols_rsq(X, Y, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X, Y)
        else:
            values = keys_core_pd_v2(load_keys, X, Y)

        def rolling_ols(df, params):
            try:
                endog = df.iloc[:, 0]
                exog = sm.add_constant(df.iloc[:, 1])

                mod = RollingOLS(endog=endog,
                                 exog=exog,
                                 window=params,
                                 min_nobs=params,
                                 missing='skip')

                model_res = mod.fit(params_only=False)

                return model_res.rsquared_adj
            except:
                ret_df = df.copy()
                ret_df.iloc[:, :] = np.nan

                return ret_df['f1']

        rsquared_adj = values.groupby(as_index=False,
                                      level='stock_code',
                                      sort=False).apply(rolling_ols, params)

        # TODO - 将inf设置为nan
        rsquared_adj[np.isinf(rsquared_adj)] = np.nan

        del load_keys

        return rsquared_adj.values

    return gp_ts_cw_ols_rsq

# 11
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ols_beta(keys, key_by_path=True):
    def gp_ts_cw_ols_beta(X, Y, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X, Y)
        else:
            values = keys_core_pd_v2(load_keys, X, Y)

        def rolling_ols(df, params):
            try:
                endog = df.iloc[:, 0]
                exog = sm.add_constant(df.iloc[:, 1])

                mod = RollingOLS(endog=endog,
                                 exog=exog,
                                 window=params,
                                 min_nobs=params,
                                 missing='skip')

                model_res = mod.fit(params_only=True)

                return model_res.params['f2']
            except:
                ret_df = df.copy()
                ret_df.iloc[:, :] = np.nan

                return ret_df['f1']

        params_beta = values.groupby(as_index=False,
                                      level='stock_code',
                                      sort=False).apply(rolling_ols, params)

        # TODO - 将inf设置为nan
        params_beta[np.isinf(params_beta)] = np.nan

        del load_keys

        return params_beta.values

    return gp_ts_cw_ols_beta

# 12
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ols_const(keys, key_by_path=True):
    def gp_ts_cw_ols_const(X, Y, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X, Y)
        else:
            values = keys_core_pd_v2(load_keys, X, Y)

        def rolling_ols(df, params):
            try:
                endog = df.iloc[:, 0]
                exog = sm.add_constant(df.iloc[:, 1])

                mod = RollingOLS(endog=endog,
                                 exog=exog,
                                 window=params,
                                 min_nobs=params,
                                 missing='skip')

                model_res = mod.fit(params_only=True)

                return model_res.params['const']
            except:
                ret_df = df.copy()
                ret_df.iloc[:, :] = np.nan

                return ret_df['f1']

        params_const = values.groupby(as_index=False,
                                      level='stock_code',
                                      sort=False).apply(rolling_ols, params)

        # TODO - 将inf设置为nan
        params_const[np.isinf(params_const)] = np.nan

        del load_keys

        return params_const.values

    return gp_ts_cw_ols_const

